#!/usr/bin/env python3
"""
发言榜图片生成脚本
使用Pillow库生成水平柱状图，替代原有的puppeteer方案
"""

import json
import sys
import os
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO
import base64
import argparse
from typing import List, Dict, Any, Optional, Tuple
from pathlib import Path


class RankingImageGenerator:
    def __init__(self, font_path: str = None):
        """
        初始化图片生成器
        
        Args:
            font_path: 字体文件路径，如果为None则使用当前目录下的字体
        """
        if font_path is None:
            # 尝试多个可能的字体文件路径
            current_dir = Path(__file__).parent
            possible_paths = [
                current_dir / "yuanshen.ttf",
                Path("/www/robot/xiuxian/plugins/nonebot_plugin_message_counter/yuanshen.ttf"),
                Path("/www/robot/xiuxian/yuanshen.ttf"),
                Path("./yuanshen.ttf"),
                Path("../yuanshen.ttf")
            ]
            
            self.font_path = None
            for path in possible_paths:
                if path.exists():
                    self.font_path = str(path)
                    break
            
            # 如果都找不到，使用默认路径
            if self.font_path is None:
                self.font_path = str(current_dir / "yuanshen.ttf")
        else:
            self.font_path = font_path
            
        self.default_font_size = 24
        self.title_font_size = 32
        self.user_avatar_size = 50
        self.bar_height = 50
        self.padding = 20
        
        self.default_avatar = self._create_default_avatar()
        self.fallback_backgrounds = self._load_fallback_backgrounds()
        
        # 尝试加载字体
        self.font = None
        self.title_font = None
        
        # 首先尝试加载指定的字体文件
        if Path(self.font_path).exists():
            try:
                self.font = ImageFont.truetype(self.font_path, self.default_font_size)
                self.title_font = ImageFont.truetype(self.font_path, self.title_font_size)
                print(f"成功加载字体文件: {self.font_path}")
            except (OSError, IOError) as e:
                print(f"警告: 无法加载字体文件 {self.font_path}: {e}")
        else:
            print(f"警告: 字体文件不存在: {self.font_path}")
        
        # 如果指定字体加载失败，尝试系统字体
        if self.font is None:
            print("尝试加载系统字体...")
            import platform
            system = platform.system()
            
            if system == "Windows":
                # Windows系统字体路径
                windows_fonts = [
                    "C:/Windows/Fonts/msyh.ttc",      # 微软雅黑
                    "C:/Windows/Fonts/simhei.ttf",    # 黑体
                    "C:/Windows/Fonts/simsun.ttc",    # 宋体
                    "C:/Windows/Fonts/simkai.ttf",    # 楷体
                    "msyh.ttc",                       # 系统字体名
                    "simhei.ttf",
                    "simsun.ttc"
                ]
                
                for font_path in windows_fonts:
                    try:
                        if font_path.startswith("C:/") and not Path(font_path).exists():
                            continue
                        self.font = ImageFont.truetype(font_path, self.default_font_size)
                        self.title_font = ImageFont.truetype(font_path, self.title_font_size)
                        print(f"成功加载Windows系统字体: {font_path}")
                        break
                    except (OSError, IOError):
                        continue
                        
            elif system == "Linux":
                # Linux系统字体
                linux_fonts = [
                    "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",
                    "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",
                    "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
                    "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf"
                ]
                
                for font_path in linux_fonts:
                    try:
                        if Path(font_path).exists():
                            self.font = ImageFont.truetype(font_path, self.default_font_size)
                            self.title_font = ImageFont.truetype(font_path, self.title_font_size)
                            print(f"成功加载Linux系统字体: {font_path}")
                            break
                    except (OSError, IOError):
                        continue
                        
            else:  # Mac系统
                mac_fonts = [
                    "/System/Library/Fonts/PingFang.ttc",
                    "/System/Library/Fonts/Helvetica.ttc",
                    "PingFang.ttc",
                    "Helvetica",
                    "Arial"
                ]
                
                for font_path in mac_fonts:
                    try:
                        if font_path.startswith("/") and not Path(font_path).exists():
                            continue
                        self.font = ImageFont.truetype(font_path, self.default_font_size)
                        self.title_font = ImageFont.truetype(font_path, self.title_font_size)
                        print(f"成功加载Mac系统字体: {font_path}")
                        break
                    except (OSError, IOError):
                        continue
        
        # 如果所有字体都加载失败，使用PIL默认字体（但会警告用户）
        if self.font is None:
            print("错误: 无法加载任何支持中文的字体，将使用PIL默认字体（可能无法显示中文）")
            self.font = ImageFont.load_default()
            self.title_font = ImageFont.load_default()
            
    def _load_fallback_backgrounds(self) -> List[str]:
        """加载fallbackBase64.json中的背景图片数据"""
        try:
            current_dir = Path(__file__).parent
            fallback_path = current_dir / "fallbackBase64.json"
            
            if fallback_path.exists():
                with open(fallback_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    return data if isinstance(data, list) else []
            return []
        except Exception as e:
            print(f"加载背景图片数据失败: {e}")
            return []
        
    def _create_default_avatar(self) -> Image.Image:
        """创建默认头像"""
        avatar = Image.new('RGB', (50, 50), color='#cccccc')
        draw = ImageDraw.Draw(avatar)
        draw.ellipse([5, 5, 45, 45], fill='#999999')
        return avatar

    def get_avatar_from_base64(self, base64_data: str) -> Optional[Image.Image]:
        """
        从base64数据获取头像图片
        
        Args:
            base64_data: base64编码的图片数据
            
        Returns:
            PIL Image对象或None
        """
        try:
            # 检查是否为空或None
            if not base64_data or base64_data.strip() == '':
                return None
                
            # 处理data URI格式
            if base64_data.startswith('data:image'):
                # 移除data URI前缀
                base64_data = base64_data.split(',', 1)[1]
            
            # 解码base64数据
            image_data = base64.b64decode(base64_data)
            image = Image.open(BytesIO(image_data))
            
            # 转换为RGBA格式以支持透明度
            return image.convert('RGBA')
        except Exception as e:
            print(f"警告: 无法解析头像数据: {e}")
            return None

    def get_average_color(self, image: Image.Image) -> Tuple[int, int, int]:
        """
        获取图片的平均颜色
        
        Args:
            image: PIL Image对象
            
        Returns:
            RGB颜色元组
        """
        try:
            # 缩小图片以提高计算速度
            small_image = image.resize((50, 50))
            # 转换为RGB模式
            if small_image.mode != 'RGB':
                small_image = small_image.convert('RGB')
            
            # 计算平均颜色
            pixels = list(small_image.getdata())
            r_sum = sum(pixel[0] for pixel in pixels)
            g_sum = sum(pixel[1] for pixel in pixels)
            b_sum = sum(pixel[2] for pixel in pixels)
            
            pixel_count = len(pixels)
            avg_r = r_sum // pixel_count
            avg_g = g_sum // pixel_count
            avg_b = b_sum // pixel_count
            
            return (avg_r, avg_g, avg_b)
        except Exception:
            # 如果出错，返回默认颜色
            return (100, 150, 200)

    def add_opacity_to_color(self, color: Tuple[int, int, int], opacity: float) -> Tuple[int, int, int, int]:
        """
        为RGB颜色添加透明度
        
        Args:
            color: RGB颜色元组
            opacity: 透明度 (0.0-1.0)
            
        Returns:
            RGBA颜色元组
        """
        return (*color, int(255 * opacity))
    
    def hex_to_rgb(self, hex_color: str) -> Tuple[int, int, int]:
        """将十六进制颜色转换为RGB"""
        hex_color = hex_color.replace("#", "")
        return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
    
    def rgb_to_hex(self, r: int, g: int, b: int) -> str:
        """将RGB颜色转换为十六进制"""
        return f"#{r:02x}{g:02x}{b:02x}"
    
    def rgb_to_hsl(self, r: int, g: int, b: int) -> Tuple[float, float, float]:
        """将RGB转换为HSL"""
        r, g, b = r / 255.0, g / 255.0, b / 255.0
        max_val = max(r, g, b)
        min_val = min(r, g, b)
        h, s, l = 0, 0, (max_val + min_val) / 2
        
        if max_val == min_val:
            h = s = 0  # achromatic
        else:
            d = max_val - min_val
            s = d / (2 - max_val - min_val) if l > 0.5 else d / (max_val + min_val)
            if max_val == r:
                h = (g - b) / d + (6 if g < b else 0)
            elif max_val == g:
                h = (b - r) / d + 2
            elif max_val == b:
                h = (r - g) / d + 4
            h /= 6
        
        return h, s, l
    
    def hsl_to_rgb(self, h: float, s: float, l: float) -> Tuple[int, int, int]:
        """将HSL转换为RGB"""
        def hue_to_rgb(p: float, q: float, t: float) -> float:
            if t < 0:
                t += 1
            if t > 1:
                t -= 1
            if t < 1/6:
                return p + (q - p) * 6 * t
            if t < 1/2:
                return q
            if t < 2/3:
                return p + (q - p) * (2/3 - t) * 6
            return p
        
        if s == 0:
            r = g = b = l  # achromatic
        else:
            q = l * (1 + s) if l < 0.5 else l + s - l * s
            p = 2 * l - q
            r = hue_to_rgb(p, q, h + 1/3)
            g = hue_to_rgb(p, q, h)
            b = hue_to_rgb(p, q, h - 1/3)
        
        return round(r * 255), round(g * 255), round(b * 255)
    
    def calculate_yiq_brightness(self, r: int, g: int, b: int) -> float:
        """计算YIQ亮度"""
        return (r * 299 + g * 587 + b * 114) / 1000 / 255
    
    def adjust_color_yiq(self, hex_color: str) -> str:
        """基于YIQ亮度调整颜色"""
        r, g, b = self.hex_to_rgb(hex_color)
        yiq_brightness = self.calculate_yiq_brightness(r, g, b)
        return "#000000" if yiq_brightness >= 0.8 else "#FFFFFF"
    
    def adjust_color_hsl(self, hex_color: str) -> str:
        """基于HSL调整颜色"""
        r, g, b = self.hex_to_rgb(hex_color)
        h, s, l = self.rgb_to_hsl(r, g, b)
        
        # 调整亮度和饱和度
        l = l + 0.3 if l < 0.5 else l - 0.3
        s = s + 0.3 if s < 0.5 else s - 0.3
        
        # 确保值在有效范围内
        l = max(0, min(1, l))
        s = max(0, min(1, s))
        
        r, g, b = self.hsl_to_rgb(h, s, l)
        return self.rgb_to_hex(r, g, b)
    
    def _generate_user_color(self, name: str, index: int, avatar_img: Optional[Image.Image] = None) -> Tuple[int, int, int]:
        """
        根据用户头像或用户名生成颜色，优先从头像提取平均颜色
        
        Args:
            name: 用户名
            index: 用户在列表中的索引
            avatar_img: 用户头像图片（可选）
            
        Returns:
            RGB颜色元组
        """
        # 如果有头像，优先从头像提取平均颜色（类似TS版本的逻辑）
        if avatar_img:
            try:
                return self.get_average_color(avatar_img)
            except Exception:
                # 如果提取失败，使用备用逻辑
                pass
        
        # 备用逻辑：预定义的颜色列表，确保相邻颜色有足够的对比度
        colors = [
            (255, 99, 132),   # 红色
            (54, 162, 235),   # 蓝色
            (255, 205, 86),   # 黄色
            (75, 192, 192),   # 青色
            (153, 102, 255),  # 紫色
            (255, 159, 64),   # 橙色
            (199, 199, 199),  # 灰色
            (83, 102, 255),   # 靛蓝色
            (255, 99, 255),   # 品红色
            (99, 255, 132),   # 绿色
        ]
        
        # 使用索引选择颜色，确保相邻用户颜色不同
        color_index = index % len(colors)
        base_color = colors[color_index]
        
        # 根据用户名的哈希值微调颜色，增加变化
        name_hash = hash(name) % 50 - 25  # -25 到 +25 的调整范围
        
        # 调整颜色，确保在有效范围内
        adjusted_color = tuple(
            max(30, min(255, base_color[i] + name_hash)) 
            for i in range(3)
        )
        
        return adjusted_color

    def choose_color_adjustment_method(self, hex_color: str) -> str:
        """选择颜色调整方法"""
        r, g, b = self.hex_to_rgb(hex_color)
        yiq_brightness = self.calculate_yiq_brightness(r, g, b)
        
        if 0.2 < yiq_brightness < 0.8:
            return self.adjust_color_hsl(hex_color)
        else:
            return self.adjust_color_yiq(hex_color)

    def generate_ranking_image(self, data: Dict[str, Any]) -> bytes:
        """
        生成排行榜图片
        
        Args:
            data: 包含排行榜数据的字典
            
        Returns:
            图片的字节数据
        """
        ranking_data = data.get('rankingData', [])
        rank_title = data.get('rankTitle', '发言排行榜')
        rank_time_title = data.get('rankTimeTitle', '')
        config = data.get('config', {})
        
        if not ranking_data:
            # 如果没有数据，生成一个空的图片
            img = Image.new('RGB', (800, 200), color=(245, 245, 245))
            draw = ImageDraw.Draw(img)
            draw.text((400, 100), "暂无数据", font=self.title_font, fill=(100, 100, 100), anchor="mm")
            
            buffer = BytesIO()
            img.save(buffer, format='PNG')
            return buffer.getvalue()
        
        # 计算图片尺寸
        user_count = len(ranking_data)
        max_count = max(item.get('count', 0) for item in ranking_data) or 1
        
        # 计算最长文本宽度
        max_text_width = 0
        for item in ranking_data:
            count = item.get('count', 0)
            percentage = item.get('percentage', 0)
            
            text = str(count)
            if config.get('isUserMessagePercentageVisible', True) and percentage > 0:
                if percentage < 0.01:
                    percentage_str = '<0.01'
                else:
                    percentage_str = f"{percentage:.2f}" if percentage < 1 else f"{percentage:.0f}"
                text += f" ({percentage_str}%)"
            
            bbox = self.font.getbbox(text)
            text_width = bbox[2] - bbox[0]
            max_text_width = max(max_text_width, text_width)
        
        # 计算图片尺寸
        bar_width = 700  # 进度条最大宽度
        image_width = self.user_avatar_size + bar_width + 10 + max_text_width + self.padding * 2
        image_height = user_count * self.bar_height + self.padding * 2
        
        # 如果有标题，增加高度
        title_height = 0
        if rank_title or rank_time_title:
            title_height = 80
            image_height += title_height
        
        # 创建图片 - 使用动态背景色
        # 计算所有用户的平均颜色作为背景基调
        background_colors = []
        for i, item in enumerate(ranking_data):
            name = item.get('name', '未知用户')
            avatar_base64 = item.get('avatarBase64', '')
            
            if avatar_base64:
                avatar_img = self.get_avatar_from_base64(avatar_base64)
                if avatar_img:
                    avg_color = self._generate_user_color(name, i, avatar_img)
                    background_colors.append(avg_color)
        
        # 如果有头像颜色，使用平均值作为背景；否则使用浅色背景
        if background_colors:
            avg_r = sum(color[0] for color in background_colors) // len(background_colors)
            avg_g = sum(color[1] for color in background_colors) // len(background_colors)
            avg_b = sum(color[2] for color in background_colors) // len(background_colors)
            # 让背景色更浅一些
            bg_color = (
                min(255, avg_r + 100),
                min(255, avg_g + 100),
                min(255, avg_b + 100)
            )
        else:
            bg_color = (245, 245, 245)  # 浅灰色背景
        
        img = Image.new('RGB', (image_width, image_height), color=bg_color)
        
        draw = ImageDraw.Draw(img)
        
        # 绘制标题
        current_y = self.padding
        if rank_title or rank_time_title:
            title_text = f"{rank_time_title} {rank_title}".strip()
            bbox = self.title_font.getbbox(title_text)
            title_width = bbox[2] - bbox[0]
            title_x = (image_width - title_width) // 2
            draw.text((title_x, current_y), title_text, font=self.title_font, fill=(51, 51, 51))
            current_y += title_height
        
        # 绘制排行榜条目
        for i, item in enumerate(ranking_data):
            y = current_y + i * self.bar_height
            
            # 获取用户信息
            name = item.get('name', '未知用户')
            count = item.get('count', 0)
            percentage = item.get('percentage', 0)
            avatar_base64 = item.get('avatarBase64', '')
            
            # 计算进度条宽度
            bar_progress_width = int((bar_width * count) / max_count) if max_count > 0 else 0
            
            # 获取头像和平均颜色
            avatar_img = None
            
            if avatar_base64:
                avatar_img = self.get_avatar_from_base64(avatar_base64)
                if avatar_img:
                    # 使用新的颜色生成逻辑，从头像提取平均颜色
                    avg_color = self._generate_user_color(name, i, avatar_img)
                else:
                    # 根据用户名生成颜色
                    avg_color = self._generate_user_color(name, i)
            else:
                # 根据用户名和索引生成颜色
                avg_color = self._generate_user_color(name, i)
            
            # 获取颜色的十六进制表示用于文本颜色调整
            avg_color_hex = self.rgb_to_hex(*avg_color)
            
            # 绘制头像
            avatar_x = self.padding
            if avatar_img:
                # 调整头像大小
                avatar_img = avatar_img.resize((self.user_avatar_size, self.user_avatar_size), Image.Resampling.LANCZOS)
                # 创建圆形遮罩
                mask = Image.new('L', (self.user_avatar_size, self.user_avatar_size), 0)
                mask_draw = ImageDraw.Draw(mask)
                mask_draw.ellipse((0, 0, self.user_avatar_size, self.user_avatar_size), fill=255)
                
                # 应用遮罩
                avatar_img.putalpha(mask)
                img.paste(avatar_img, (avatar_x, y), avatar_img)
            else:
                # 绘制默认头像（圆形）
                draw.ellipse((avatar_x, y, avatar_x + self.user_avatar_size, y + self.user_avatar_size), 
                           fill=avg_color, outline=(200, 200, 200))
                # 在圆形中绘制用户名首字符
                if name:
                    first_char = name[0]
                    char_bbox = self.font.getbbox(first_char)
                    char_width = char_bbox[2] - char_bbox[0]
                    char_height = char_bbox[3] - char_bbox[1]
                    char_x = avatar_x + (self.user_avatar_size - char_width) // 2
                    char_y = y + (self.user_avatar_size - char_height) // 2
                    draw.text((char_x, char_y), first_char, font=self.font, fill=(255, 255, 255))
            
            # 绘制进度条
            bar_x = avatar_x + self.user_avatar_size
            
            # 绘制进度条背景（完整宽度）
            bg_color = self.add_opacity_to_color(avg_color, 0.3)
            bg_rect = Image.new('RGBA', (bar_width, self.bar_height), bg_color)
            img.paste(bg_rect, (bar_x, y), bg_rect)
            
            # 绘制进度条前景
            if bar_progress_width > 0:
                fg_rect = Image.new('RGBA', (bar_progress_width, self.bar_height), (*avg_color, 255))
                img.paste(fg_rect, (bar_x, y), fg_rect)
            
            # 绘制用户名 - 使用颜色调整方法确保文本可读性
            name_x = bar_x + 10
            name_y = y + (self.bar_height - self.default_font_size) // 2
            text_color = self.choose_color_adjustment_method(avg_color_hex)
            text_color_rgb = self.hex_to_rgb(text_color)
            draw.text((name_x, name_y), name, font=self.font, fill=text_color_rgb)
            
            # 绘制计数和百分比
            count_text = str(count)
            if config.get('isUserMessagePercentageVisible', True) and percentage > 0:
                if percentage < 0.01:
                    percentage_str = '<0.01'
                else:
                    percentage_str = f"{percentage:.2f}" if percentage < 1 else f"{percentage:.0f}"
                count_text += f" ({percentage_str}%)"
            
            count_width = draw.textlength(count_text, font=self.font)
            count_x = bar_x + bar_width + 10
            count_y = y + (self.bar_height - self.default_font_size) // 2
            
            # 如果文本超出边界，使用调整后的颜色在柱状图内显示
            if count_x + count_width > image_width - 10:
                draw.text((bar_x + bar_width - count_width - 10, count_y), 
                         count_text, fill=text_color_rgb, font=self.font)
            else:
                draw.text((count_x, count_y), count_text, 
                         fill=(51, 51, 51), font=self.font)
        
        # 保存图片到字节流
        buffer = BytesIO()
        img.save(buffer, format='PNG', quality=95)
        return buffer.getvalue()


def main():
    """主函数，处理命令行参数"""
    if len(sys.argv) != 2:
        print("用法: python generate_ranking_image.py <json_data>")
        sys.exit(1)
    
    try:
        # 解析JSON数据
        json_data = sys.argv[1]
        data = json.loads(json_data)
        
        # 获取字体路径
        font_path = "yuanshen.ttf"
        if not os.path.exists(font_path):
            # 尝试在当前目录查找
            current_dir = os.path.dirname(os.path.abspath(__file__))
            font_path = os.path.join(current_dir, "yuanshen.ttf")
        
        # 创建图片生成器
        generator = RankingImageGenerator(font_path)
        
        # 生成图片
        image_data = generator.generate_ranking_image(data)
        
        # 输出base64编码的图片数据
        base64_data = base64.b64encode(image_data).decode('utf-8')
        print(base64_data)
        
    except json.JSONDecodeError as e:
        print(f"JSON解析错误: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"生成图片时发生错误: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()